{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "abd165c0-216b-475d-b939-c8b7da507b51",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import sys\n",
    "import os\n",
    "import json\n",
    "import cv2\n",
    "import uuid\n",
    "sys.path.append('../../../src')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "81558f18-d8cf-4a61-b25a-407ef8e19430",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sr_od.utils import os_utils\n",
    "from sr_od.devtools import dataset_utils\n",
    "from sr_od.devtools import label_studio_utils\n",
    "from sr_od.devtools import ultralytics_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "24e42133-abe2-4274-af18-b3a8dada4a25",
   "metadata": {},
   "outputs": [],
   "source": [
    "ultralytics_utils.init_ultralytics_settings()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dda76c31-9545-4660-879b-b705a96bf073",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ultralytics import YOLO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "eff855ef-5c8b-4f29-bc69-1dea9c9a6f42",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = YOLO(ultralytics_utils.get_train_model_path('full-test-v0', 'train', 'best'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6c7dabcc-7d11-4196-843e-c7f952332b98",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ultralytics YOLOv8.1.44  Python-3.11.9 torch-2.2.2+cu121 CPU (AMD Ryzen 9 5950X 16-Core Processor)\n",
      "Model summary (fused): 168 layers, 3015593 parameters, 0 gradients, 8.1 GFLOPs\n",
      "\n",
      "\u001B[34m\u001B[1mPyTorch:\u001B[0m starting from 'F:\\code\\workspace\\StarRail-YOLO\\ultralytics\\runs\\full-test-v0\\train\\weights\\best.pt' with input shape (1, 3, 384, 640) BCHW and output shape(s) (1, 55, 5040) (6.0 MB)\n",
      "\n",
      "\u001B[34m\u001B[1mONNX:\u001B[0m starting export with onnx 1.16.0 opset 17...\n",
      "\u001B[34m\u001B[1mONNX:\u001B[0m export success  1.0s, saved as 'F:\\code\\workspace\\StarRail-YOLO\\ultralytics\\runs\\full-test-v0\\train\\weights\\best.onnx' (11.6 MB)\n",
      "\n",
      "Export complete (2.7s)\n",
      "Results saved to \u001B[1mF:\\code\\workspace\\StarRail-YOLO\\ultralytics\\runs\\full-test-v0\\train\\weights\u001B[0m\n",
      "Predict:         yolo predict task=detect model=F:\\code\\workspace\\StarRail-YOLO\\ultralytics\\runs\\full-test-v0\\train\\weights\\best.onnx imgsz=384,640  \n",
      "Validate:        yolo val task=detect model=F:\\code\\workspace\\StarRail-YOLO\\ultralytics\\runs\\full-test-v0\\train\\weights\\best.onnx imgsz=384,640 data=F:\\code\\workspace\\StarRail-YOLO\\ultralytics\\datasets\\full-test-v0\\dataset.yaml  WARNING  non-PyTorch val requires square images, 'imgsz=[384, 640]' will not work. Use export 'imgsz=640' if val is required.\n",
      "Visualize:       https://netron.app\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'F:\\\\code\\\\workspace\\\\StarRail-YOLO\\\\ultralytics\\\\runs\\\\full-test-v0\\\\train\\\\weights\\\\best.onnx'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.export(format='onnx', imgsz=(384, 640))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d86865f0-40ac-43a7-861c-12fbccf66076",
   "metadata": {},
   "outputs": [],
   "source": "from sr_od.detector import YOLOv8"
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b5576e09-4bc1-4891-a774-3afd7b59d972",
   "metadata": {},
   "outputs": [],
   "source": [
    "onnx_path = 'F:\\\\code\\\\workspace\\\\StarRail-YOLO\\\\ultralytics\\\\runs\\\\xuzu-test\\\\train\\\\weights\\\\best.onnx'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "cdc337cb-ea56-45e3-ae4c-6d99efcec77d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = YOLOv8(onnx_path, cuda=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c90a5fb5-504a-44c5-9e65-e37a307bd58f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4a7128be-3b83-403d-bcbc-bdb76d97820f",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name = 'xuzu-test'\n",
    "img_dir = dataset_utils.get_dataset_images_dir(dataset_name)\n",
    "img_list = os.listdir(img_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ed8b6409-6ded-418d-b05f-c8d14f3e3160",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "random.seed()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "42469970-49b6-48e6-bf6c-1cf050600b1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_txt_path = os.path.join(ultralytics_utils.get_dataset_dir(dataset_name), 'autosplit_test.txt')\n",
    "with open(test_txt_path, 'r') as file:\n",
    "    lines = file.readlines()\n",
    "    test_case_ids = [line.strip()[9:] for line in lines]  # 去除每行末尾的换行符"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4488fe3e-1624-4bfe-8287-e859dc008616",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79df7f99-5092-4c6b-a3e2-30ef34866b09",
   "metadata": {},
   "outputs": [],
   "source": [
    "random_img = os.path.join(img_dir, random.choice(test_case_ids))\n",
    "img = cv2.imread(random_img)\n",
    "t1 = time.time()\n",
    "model(img)\n",
    "print(time.time() - t1)\n",
    "result_img = model.draw_detections(img)\n",
    "plt.imshow(cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab9b343f-07f8-44b6-901b-fbc3551cdd63",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96c69cb6-8001-4e10-8c83-eecdc73462ef",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
